#!/usr/bin/env python

import roslib; roslib.load_manifest('team4')
import rospy
from pr2_tasks.tasks import Tasks
from pr2_tasks.pickplace import PickPlace
from pr2_tasks.tableware_detection import TablewareDetection
from pr2_tasks.pickplace_definitions import PlaceGoal
from pr2_python import arm_planner, transform_listener
from arm_mover import ArmMover
from object_manipulation_msgs.msg import Grasp, ClusterBoundingBox
from pr2_controllers_msgs.msg import JointTrajectoryAction
from geometry_msgs.msg import PointStamped, PoseStamped
from trajectory_msgs.msg import JointTrajectoryPoint, JointTrajectory
from tf.transformations import *
from sensor_msgs.msg import PointCloud

class Tray():
   def __init__(self):
      self.tray_location = PointStamped()
      self.last_grasps = None
      self.tray = None
      self.b_right = None
      self.b_left = None
      self.l_arm_mover = ArmMover('left_arm')
      self.r_arm_mover = ArmMover('right_arm')
      self.obj_bounding_box = None
      self.listener = transform_listener.get_transform_listener()

   def find_tray(self):
      det = TablewareDetection()
      head_point = PointStamped()
      head_point.header.frame_id = '/torso_lift_link'
      head_point.point.x = 0.6
      res = det.detect_objects()
      
      for pg in res.pickup_goals:
         if pg.label == 'graspable': # limit this to unrecognized objects, in our context trays
            self.obj_bounding_box = det._bounding_box_srv(pg.target.cluster) 
            self.tray = pg
      return

   def splitCluster(self):
      self.b_right = PointCloud()
      self.b_right.header = self.tray.target.cluster.header
      self.b_left = PointCloud()
      self.b_left.header = self.tray.target.cluster.header
      x_max = 0.0
      x_min = 100.0
      y_max = -100.0
      y_min = 100.0

      rospy.loginfo("Preparing to filter cluster of size {0}".format(len(self.tray.target.cluster.points)))
      # Find cluster stats
      for pt in self.tray.target.cluster.points:
         temp = PointStamped()
         temp.header = self.tray.target.cluster.header
         temp.point  = pt
         temp = self.listener.transformPoint("base_link", temp)
         if temp.point.x > x_max: 
            x_max = temp.point.x
         elif temp.point.x < x_min:
            x_min = temp.point.x
         if temp.point.y > y_max:
            y_max = temp.point.y
         elif temp.point.y < y_min:
            y_min = temp.point.y      
         pt = temp.point
          
      x_range = x_max - x_min
      x_center = (x_max+x_min)/2.0
      y_range = y_max - y_min
      y_center = (y_min+y_max)/2.0
      
      # Filter cluster
      for pt in self.tray.target.cluster.points:
         if pt.y >= y_center+(y_range/4.0) and ((pt.x <= x_center + (x_range/6.0)) and (pt.x >= x_center - (x_range/6.0))):
            self.b_left.points.append(pt)
         elif pt.y <= y_center-(y_range/4.0) and ((pt.x <= x_center + (x_range/6.0)) and (pt.x >= x_center - (x_range/6.0))): 
            self.b_right.points.append(pt)

      rospy.loginfo("Right cluster size: {0}".format(len(self.b_right.points)))
      rospy.loginfo("Left cluster size: {0}".format(len(self.b_left.points)))
      

   def gripClusters(self):
      pickplace = PickPlace()
      pg = self.tray
      pg.target.cluster = self.b_right
      pg.arm_name = 'right_arm'
      pg.lift.min_distance = 0.005
      pg.lift.desired_distance = 0.01#
      pickplace.pickup(pg)

      pg.target.cluster = self.b_left
      pg.arm_name = 'left_arm'
      pg.lift.min_distance = 0.005
      pg.lift.desired_distance = 0.01
      pickplace.pickup(pg)

 
   def getArmIK(self, pose=PoseStamped(), arm='right_arm'):
      arm_planner=[]
      if arm == 'right_arm':
	      arm_planner = self.r_arm_mover._arm_planner
      elif arm == 'left_arm':
	      arm_planner = self.l_arm_mover._arm_planner
      ik_sol = arm_planner.get_ik(pose)
      return ik_sol

   def moveTray(self, z_dist):
      rospy.loginfo("Waiting on transform")
      transform_listener.get_transform_as_tr('/base_link', '/r_wrist_roll_link')
      rospy.loginfo("Found transform")

      # Calculate new pose relative to current pose
      l_ps = PoseStamped()
      l_ps.header.frame_id = "/base_link"
      l_ps.header.stamp = rospy.Time(0)
      l_ps.pose.orientation.w = 1.0
      r_ps = PoseStamped()
      r_ps.header.frame_id = "/base_link"
      r_ps.header.stamp = rospy.Time(0)
      r_ps.pose.orientation.w = 1.0
      new_l_ps = transform_listener.transform_pose("/base_link", '/l_wrist_roll_link', l_ps.pose)
      new_r_ps = transform_listener.transform_pose("/base_link", '/r_wrist_roll_link', r_ps.pose)
      new_l_ps.position.z += z_dist
      new_r_ps.position.z += z_dist
      l_ps.pose = new_l_ps
      r_ps.pose = new_r_ps

      # Create new plans
      r_arm_planner = self.r_arm_mover._arm_planner
      l_arm_planner = self.l_arm_mover._arm_planner
      r_traj = r_arm_planner.plan_interpolated_ik(r_ps, collision_aware=False)
      l_traj = l_arm_planner.plan_interpolated_ik(l_ps, collision_aware=False)

      # Splice plans into one trajectory
      spliced = r_traj
      spliced.joint_names = r_traj.joint_names + l_traj.joint_names
      m = min(len(r_traj.points), len(l_traj.points))
      for i in range(0,m):
         rospy.loginfo("Splicing point {0}".format(i))
         spliced.points[i].positions = r_traj.points[i].positions + l_traj.points[i].positions
         spliced.points[i].velocities = r_traj.points[i].velocities + l_traj.points[i].velocities
      rospy.loginfo("Executing spliced trajectory")

      # Exec two arm trajectory
      self.l_arm_mover.execute_two_arm_trajectory(spliced)
      rospy.loginfo("Executed spliced trajectory")

      
   def pick_up_tray(self):
      self.find_tray()
      self.splitCluster()
      self.gripClusters()
      self.moveTray(0.1)

   def put_down_tray(self):
      self.moveTray(-0.1)
      # Store location?

   def get_tray_out_of_the_way(self):
      self.moveTray(-0.5)

 #
#if __name__ == "__main__":
   #rospy.init_node("move_both_arms_test")
   #test = Tray()
     #
   #test.pick_up_tray()
   #
   #rospy.sleep(2.0)
   #test.put_down_tray()
     
